from dataclasses import dataclass, field
from typing import Dict, List
from .common import Simulator, CostCalculator
from abc import ABC, abstractmethod
from .base_simulator import (
    ColumnParallelLinearSimulator,
    RowParallelLinearSimulator,
    IdentityOpSimulator,
    DotProductAttentionSimulator,
    LNImplSimulator,
    DropoutSimulator,
    FusedScaleMaskSoftmaxSimulator,
    IdentityFuncOpSimulator,
)

@dataclass
class SelfAttentionSimulator(CostCalculator):
    # 计算QKV矩阵
    linear_qkv: ColumnParallelLinearSimulator = None
    # pre_cross_attn_layernorm
    pre_cross_attn_layernorm: IdentityOpSimulator = None
    # Attention product
    core_attention: DotProductAttentionSimulator = None
    # core_attention 点积计算后soft_max
    scale_mask_softmax: FusedScaleMaskSoftmaxSimulator = None
    # core_attention 点积后的droupout
    attention_dropout: DropoutSimulator = None

    # 投影到多维
    linear_proj: RowParallelLinearSimulator = None
    q_layernorm: IdentityOpSimulator = None
    k_layernorm: IdentityOpSimulator = None

    cross_attention: IdentityOpSimulator = None

class GPTSelfAttentionAbstructBuilder(ABC):
    @abstractmethod
    def build_linear_qkv(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_pre_cross_attn_layernorm(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_core_attention(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_scale_mask_softmax(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_attention_dropout(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_linear_proj(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_q_layernorm(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_k_layernorm(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_cross_attention(self, args: Dict[str, object]): ...

class LlamaSelfAttentionAbstructBuilder(ABC):
    pass


class SelfAttentionAbstructBuilder(GPTSelfAttentionAbstructBuilder,LlamaSelfAttentionAbstructBuilder):
    pass